// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#pragma once

#include <pollux/serializers/compact_row_serializer.h>
#include <pollux/vector/complex_vector.h>
#include <pollux/vector/vector_stream.h>

namespace kumo::pollux::serializer {
    using TRowSize = uint32_t;

    namespace detail {
        struct RowGroupHeader {
            int32_t uncompressedSize;
            int32_t compressedSize;
            bool compressed;

            static RowGroupHeader read(ByteInputStream *source) {
                RowGroupHeader header;
                header.uncompressedSize = source->read<int32_t>();
                header.compressedSize = source->read<int32_t>();
                header.compressed = source->read<char>();

                POLLUX_CHECK_GE(header.uncompressedSize, 0);
                POLLUX_CHECK_GE(header.compressedSize, 0);

                return header;
            }

            void write(OutputStream *out) {
                out->write(reinterpret_cast<char *>(&uncompressedSize), sizeof(int32_t));
                out->write(reinterpret_cast<char *>(&compressedSize), sizeof(int32_t));
                const char writeValue = compressed ? 1 : 0;
                out->write(reinterpret_cast<const char *>(&writeValue), sizeof(char));
            }

            void write(char *out) {
                ::memcpy(out, reinterpret_cast<char *>(&uncompressedSize), sizeof(int32_t));
                ::memcpy(
                    out + sizeof(int32_t),
                    reinterpret_cast<char *>(&compressedSize),
                    sizeof(int32_t));
                const char writeValue = compressed ? 1 : 0;
                ::memcpy(
                    out + sizeof(int32_t) * 2,
                    reinterpret_cast<const char *>(&writeValue),
                    sizeof(char));
            }

            static size_t size() {
                return sizeof(int32_t) * 2 + sizeof(char);
            }

            std::string debugString() const {
                return fmt::format(
                    "uncompressedSize: {}, compressedSize: {}, compressed: {}",
                    succinctBytes(uncompressedSize),
                    succinctBytes(compressedSize),
                    compressed);
            }
        };
    } // namespace detail

    template<class Serializer>
    class RowSerializer : public IterativeVectorSerializer {
    public:
        RowSerializer(memory::MemoryPool *pool, const VectorSerde::Options *options)
            : pool_(pool),
              options_(options == nullptr ? VectorSerde::Options() : *options),
              codec_(common::compressionKindToCodec(options_.compressionKind)) {
        }

        void append(
            const RowVectorPtr &vector,
            const melon::Range<const IndexRange *> &ranges,
            Scratch & /*scratch*/) override {
            size_t totalSize = 0;
            const auto totalRows = std::accumulate(
                ranges.begin(),
                ranges.end(),
                0,
                [](vector_size_t sum, const auto &range) { return sum + range.size; });

            if (totalRows == 0) {
                return;
            }

            Serializer row(vector);
            std::vector<vector_size_t> rowSize(totalRows);
            if (auto fixedRowSize =
                    Serializer::fixedRowSize(as_row_type(vector->type()))) {
                totalSize += (fixedRowSize.value() + sizeof(TRowSize)) * totalRows;
                std::fill(rowSize.begin(), rowSize.end(), fixedRowSize.value());
            } else {
                vector_size_t index = 0;
                for (const auto &range: ranges) {
                    for (auto i = 0; i < range.size; ++i, ++index) {
                        rowSize[index] = row.rowSize(range.begin + i);
                        totalSize += rowSize[index] + sizeof(TRowSize);
                    }
                }
            }

            if (totalSize == 0) {
                return;
            }

            BufferPtr buffer = AlignedBuffer::allocate<char>(totalSize, pool_, 0);
            auto *rawBuffer = buffer->asMutable<char>();
            buffers_.push_back(std::move(buffer));

            serializeRanges(row, ranges, rawBuffer, rowSize);
        }

        void append(
            const Serializer &compactRow,
            const melon::Range<const vector_size_t *> &rows,
            const std::vector<vector_size_t> &sizes) override {
            size_t totalSize = 0;
            for (const auto row: rows) {
                totalSize += sizes[row];
            }
            if (totalSize == 0) {
                return;
            }

            BufferPtr buffer = AlignedBuffer::allocate<char>(totalSize, pool_, 0);
            auto *rawBuffer = buffer->asMutable<char>();
            buffers_.push_back(std::move(buffer));

            size_t offset = 0;
            for (auto &row: rows) {
                // Write row data.
                const TRowSize size =
                        compactRow.serialize(row, rawBuffer + offset + sizeof(TRowSize));

                // Write raw size. Needs to be in big endian order.
                *(TRowSize *) (rawBuffer + offset) = melon::Endian::big(size);
                offset += sizeof(TRowSize) + size;
            }
        }

        size_t maxSerializedSize() const override {
            const auto size = uncompressedSize();
            if (!needCompression()) {
                return detail::RowGroupHeader::size() + size;
            }
            POLLUX_CHECK_LE(
                size,
                codec_->maxUncompressedLength(),
                "UncompressedSize exceeds limit");
            return detail::RowGroupHeader::size() + codec_->maxCompressedLength(size);
        }

        /// The serialization format is | uncompressedSize | compressedSize |
        /// compressed | data.
        void flush(OutputStream *stream) override {
            constexpr int32_t kMaxCompressionAttemptsToSkip = 30;
            const auto size = uncompressedSize();
            if (!needCompression()) {
                flushUncompressed(size, stream);
            } else if (numCompressionToSkip_ > 0) {
                flushUncompressed(size, stream);
                stats_.compressionSkippedBytes += size;
                --numCompressionToSkip_;
                ++stats_.numCompressionSkipped;
            } else {
                // Compress the buffer if satisfied condition.
                const auto toCompress = toIOBuf(buffers_);
                const auto compressedBuffer = codec_->compress(toCompress.get());
                const int32_t compressedSize = compressedBuffer->length();
                stats_.compressionInputBytes += size;
                stats_.compressedBytes += compressedSize;
                if (compressedSize > options_.minCompressionRatio * size) {
                    // Skip this compression.
                    numCompressionToSkip_ = std::min<int64_t>(
                        kMaxCompressionAttemptsToSkip, 1 + stats_.numCompressionSkipped);
                    flushUncompressed(size, stream);
                } else {
                    // Do the compression.
                    detail::RowGroupHeader header = {size, compressedSize, true};
                    header.write(stream);
                    for (auto range: *compressedBuffer) {
                        stream->write(
                            reinterpret_cast<const char *>(range.data()), range.size());
                    }
                }
            }

            buffers_.clear();
        }

        std::unordered_map<std::string, RuntimeCounter> runtimeStats() override {
            std::unordered_map<std::string, RuntimeCounter> map;
            if (stats_.compressionInputBytes != 0) {
                map.emplace(
                    kCompressionInputBytes,
                    RuntimeCounter(
                        stats_.compressionInputBytes, RuntimeCounter::Unit::kBytes));
            }
            if (stats_.compressedBytes != 0) {
                map.emplace(
                    kCompressedBytes,
                    RuntimeCounter(stats_.compressedBytes, RuntimeCounter::Unit::kBytes));
            }
            if (stats_.compressionSkippedBytes != 0) {
                map.emplace(
                    kCompressionSkippedBytes,
                    RuntimeCounter(
                        stats_.compressionSkippedBytes, RuntimeCounter::Unit::kBytes));
            }
            return map;
        }

        void clear() override {
        }

    protected:
        virtual void serializeRanges(
            const Serializer &rowSerializer,
            const melon::Range<const IndexRange *> &ranges,
            char *rawBuffer,
            const std::vector<vector_size_t> & /*rowSize*/) {
            size_t offset = 0;
            for (auto &range: ranges) {
                for (auto row = range.begin; row < range.begin + range.size; ++row) {
                    // Write row data.
                    TRowSize size =
                            rowSerializer.serialize(row, rawBuffer + offset + sizeof(TRowSize));

                    // Write raw size. Needs to be in big endian order.
                    *(TRowSize *) (rawBuffer + offset) = melon::Endian::big(size);
                    offset += sizeof(TRowSize) + size;
                }
            }
        }

        memory::MemoryPool *const pool_;
        std::vector<BufferPtr> buffers_;

    private:
        std::unique_ptr<melon::IOBuf> toIOBuf(const std::vector<BufferPtr> &buffers) {
            std::unique_ptr<melon::IOBuf> iobuf;
            for (const auto &buffer: buffers) {
                auto newBuf =
                        melon::IOBuf::wrapBuffer(buffer->asMutable<char>(), buffer->size());
                if (iobuf) {
                    iobuf->prev()->appendChain(std::move(newBuf));
                } else {
                    iobuf = std::move(newBuf);
                }
            }
            return iobuf;
        }

        int32_t uncompressedSize() const {
            int32_t totalSize = 0;
            for (const auto &buffer: buffers_) {
                totalSize += buffer->size();
            }
            return totalSize;
        }

        bool needCompression() const {
            return codec_->type() != melon::compression::CodecType::NO_COMPRESSION;
        }

        void flushUncompressed(int32_t size, OutputStream *stream) {
            detail::RowGroupHeader header = {size, size, false};
            header.write(stream);
            for (const auto &buffer: buffers_) {
                stream->write(buffer->template asMutable<char>(), buffer->size());
            }
        }

        const VectorSerde::Options options_;
        const std::unique_ptr<melon::compression::Codec> codec_;
        // Count of forthcoming compressions to skip.
        int32_t numCompressionToSkip_{0};
        CompressionStats stats_;
    };

    /*
     * RowIterator is an iterator to read deserialize rows from an uncompressed
     * input source. It provides the ability to iterate over the stream by
     * retrieving a serialized buffer containing one row in each iteration.
     *
     * Usage:
     * RowIterator iterator(source, endOffset);
     * while (iterator.hasNext()) {
     *   auto next = iterator.next();
     *   // Process the data in next
     * }
     */
    class RowIterator {
    public:
        virtual ~RowIterator() = default;

        virtual bool hasNext() const = 0;

        virtual std::unique_ptr<std::string> next() = 0;

    protected:
        RowIterator(ByteInputStream *source, size_t endOffset)
            : source_(source), endOffset_(endOffset) {
        }

        ByteInputStream *const source_;
        const size_t endOffset_;
    };

    class RowIteratorImpl : public RowIterator {
    public:
        RowIteratorImpl(ByteInputStream *source, size_t endOffset)
            : RowIterator(source, endOffset) {
            POLLUX_CHECK_NOT_NULL(source, "Source cannot be null");
        }

        bool hasNext() const override {
            return source_->tellp() < endOffset_;
        }

        std::unique_ptr<std::string> next() override {
            const auto rowSize = readRowSize();
            auto serializedBuffer = std::make_unique<std::string>();
            serializedBuffer->reserve(rowSize);

            const auto row = source_->nextView(rowSize);
            serializedBuffer->append(row.data(), row.size());
            // If we couldn't read the entire row at once, we need to concatenate it
            // in a different buffer.
            if (serializedBuffer->size() < rowSize) {
                concatenatePartialRow(source_, rowSize, *serializedBuffer);
            }

            POLLUX_CHECK_EQ(serializedBuffer->size(), rowSize);
            return serializedBuffer;
        }

    private:
        TRowSize readRowSize() {
            return melon::Endian::big(source_->read<TRowSize>());
        }

        // Read from the stream until the full row is concatenated.
        static void concatenatePartialRow(
            ByteInputStream *source,
            TRowSize rowSize,
            std::string &rowBuffer) {
            while (rowBuffer.size() < rowSize) {
                const std::string_view rowFragment =
                        source->nextView(rowSize - rowBuffer.size());
                POLLUX_CHECK_GT(
                    rowFragment.size(),
                    0,
                    "Unable to read full serialized row. Needed {} but read {} bytes.",
                    rowSize - rowBuffer.size(),
                    rowFragment.size());
                rowBuffer.append(rowFragment.data(), rowFragment.size());
            }
        }
    };

    template<typename SerializeView>
    class RowDeserializer {
    public:
        template<typename RowIterator>
        static void deserialize(
            ByteInputStream *source,
            std::vector<SerializeView> &serializedRows,
            std::vector<std::unique_ptr<std::string> > &serializedBuffers,
            const VectorSerde::Options *options) {
            const auto compressionKind = options == nullptr
                                             ? VectorSerde::Options().compressionKind
                                             : options->compressionKind;
            while (!source->atEnd()) {
                std::unique_ptr<melon::IOBuf> uncompressedBuf;
                const auto header = detail::RowGroupHeader::read(source);
                if (header.compressed) {
                    POLLUX_DCHECK_NE(
                        compressionKind, common::CompressionKind::CompressionKind_NONE);
                    auto compressBuf = melon::IOBuf::create(header.compressedSize);
                    source->readBytes(compressBuf->writableData(), header.compressedSize);
                    compressBuf->append(header.compressedSize);

                    // Process chained uncompressed results IOBufs.
                    const auto codec = common::compressionKindToCodec(compressionKind);
                    uncompressedBuf =
                            codec->uncompress(compressBuf.get(), header.uncompressedSize);
                }

                std::unique_ptr<ByteInputStream> uncompressedStream;
                ByteInputStream *uncompressedSource{nullptr};
                if (uncompressedBuf == nullptr) {
                    uncompressedSource = source;
                } else {
                    uncompressedStream = std::make_unique<BufferInputStream>(
                        byteRangesFromIOBuf(uncompressedBuf.get()));
                    uncompressedSource = uncompressedStream.get();
                }
                const std::streampos initialSize = uncompressedSource->tellp();
                RowIterator rowIterator(
                    uncompressedSource, header.uncompressedSize + initialSize);
                while (rowIterator.hasNext()) {
                    serializedBuffers.emplace_back(std::move(rowIterator.next()));
                    serializedRows.push_back(std::string_view(
                        serializedBuffers.back()->data(),
                        serializedBuffers.back()->size()));
                }
            }
        }
    };
} // namespace kumo::pollux::serializer
